package com.hzw.saas.service.hpc.util;

import cn.hutool.core.date.DateTime;
import cn.hutool.core.date.DateUnit;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.text.StrFormatter;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.ssh.JschUtil;
import com.hzw.saas.api.hpc.bo.HpcAccountBo;
import com.hzw.saas.common.util.Arith;
import com.hzw.saas.service.hpc.config.SshConfig;
import com.hzw.saas.service.hpc.data.HACDetail;
import com.hzw.saas.service.hpc.exception.PullHpcDetailsException;
import com.hzw.saas.service.hpc.service.SshHub;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;

/**
 * <p>
 *
 * </p>
 *
 * @author sonam
 * @since 2021/4/28 10:17
 */
@Slf4j
public class HpcUtils {

    /**
     * ssh拉取HPC用户明细
     * TODO optimize
     *
     * @param hpcAccountPoolBo
     * @param startDateStr
     * @param endDateStr
     * @return
     * @throws IOException
     * @throws JSchException
     */
    public static List<HACDetail> pullHpcAccountDetails(HpcAccountBo hpcAccountPoolBo, String startDateStr, String endDateStr) throws PullHpcDetailsException {
        String hpcAccount = hpcAccountPoolBo.getHpcAccount();
        String hpcPassword = hpcAccountPoolBo.getHpcPassword();
        Integer port = hpcAccountPoolBo.getPort();
        String host = hpcAccountPoolBo.getHost();
        String[] strings = hpcAccount.split("@");
        String simpleHpcAccount = hpcAccount;
        String cluster = StrUtil.EMPTY;
        if (strings.length > 1) {
            simpleHpcAccount = strings[0];
            cluster = strings[1];
        }

        SshConfig sshConfig = new SshConfig();
        sshConfig.setHost(host);
        sshConfig.setPort(port);
        sshConfig.setUsername(hpcAccount);
        sshConfig.setPassword(hpcPassword);

        SshHub sshHub = new SshHub();
        sshHub.setConfig(sshConfig);
        try {
            // 获取正在运行的任务命令 sacct -u sc80264 -X -T
            String command = "";
            if (StrUtil.isAllBlank(startDateStr, endDateStr)) {
                command = StrFormatter.format("sacct -u {} -X -T --field=jobid,jobname,user,start,end,state,reqcpus", simpleHpcAccount);
            } else {
                command = StrFormatter.format("sacct -u {} -S {} -E {} -X -T --field=jobid,jobname,user,start,end,state,reqcpus", simpleHpcAccount, startDateStr, endDateStr);
            }
            sshHub.execute(command);
        } catch (Exception e) {
            throw new PullHpcDetailsException(e.getMessage(), e);
        }
        List<HACDetail> hapDetails = new ArrayList<>();
        List<String> results = sshHub.getResult();
        for (int i = 2; i < results.size(); i++) {
            String result = results.get(i);
            String[] dataArray = result.split("\\s+");
            String jodIdStr = dataArray[0];
            String jobNameStr = dataArray[1];
            String userStr = dataArray[2];
            String startStr = dataArray[3].replace("T", " ");
            String endStr = dataArray[4].replace("T", " ");
            String stateStr = dataArray[5];
            String reqCpusStr = dataArray[6];

            DateTime startTime = DateUtil.parse(startStr);
            DateTime endTime = DateUtil.parse(endStr);

            Long sumTimeS = DateUtil.between(startTime, endTime, DateUnit.SECOND);
            int allocCpus = Integer.parseInt(reqCpusStr);
            HACDetail detail = new HACDetail().setStartTime(startTime).setEndTime(endTime).setJobId(jodIdStr).setSubmitTime(startTime).setAllocCpus(allocCpus).setCputimeRaw(Arith.mul(allocCpus, sumTimeS)) // CPU核数 * 运行时长秒
                .setUser(userStr).setJobName(jobNameStr).setCluster(cluster).setStatus(stateStr);
            hapDetails.add(detail);
        }

        return hapDetails;
    }

    public static List<String> pullJobIdByStatus(Session session, String hpcAccount, String state) {
        String command = StrFormatter.format("sacct -u {} -S 2020-01-01 -E {} -X -T --field=jobid,state", hpcAccount, DateUtil.formatDate(DateUtil.offsetDay(new Date(), 1)));
        String jobStr = JschUtil.exec(session, command, Charset.defaultCharset());
        String[] jobs = jobStr.split(StrUtil.LF);
        List<String> jobIds = new ArrayList<>();
        for (String job : jobs) {
            if (job.contains(state)) {
                jobIds.add(job.split(StrUtil.SPACE)[0]);
            }
        }
        return jobIds;
    }

//    public static void main(String[] args) {
//        Session session = JschUtil.openSession("60.216.39.134", 22022, "sonam-03", "Y4hAGyd-30xHjS_eMEUY");
//        System.out.println(pullJobIdByStatus(session, "sonam-03", "PENDING"));
//    }
}
